open Core_kernel
open Currency
open Signature_lib
open Mina_base
module Ledger = Mina_ledger.Ledger
module Root_ledger = Mina_ledger.Root
module Intf = Intf
module Ledger_transfer_mask =
  Mina_ledger.Ledger_transfer.Make (Ledger) (Ledger.Any_ledger.M)
module Ledger_transfer_any =
  Mina_ledger.Ledger_transfer.Make (Ledger.Any_ledger.M) (Ledger.Any_ledger.M)

let account_with_timing account_id balance (timing : Intf.Timing.t) =
  match timing with
  | Untimed ->
      Account.create account_id balance
  | Timed t ->
      let initial_minimum_balance =
        Currency.Balance.of_nanomina_int_exn t.initial_minimum_balance
      in
      let cliff_time =
        Mina_numbers.Global_slot_since_genesis.of_int t.cliff_time
      in
      let cliff_amount = Currency.Amount.of_nanomina_int_exn t.cliff_amount in
      let vesting_increment =
        Currency.Amount.of_nanomina_int_exn t.vesting_increment
      in
      let vesting_period =
        Mina_numbers.Global_slot_span.of_int t.vesting_period
      in
      Account.create_timed account_id balance ~initial_minimum_balance
        ~cliff_time ~cliff_amount ~vesting_period ~vesting_increment
      |> Or_error.ok_exn

module Private_accounts (Accounts : Intf.Private_accounts.S) = struct
  include Accounts

  let accounts =
    let open Lazy.Let_syntax in
    let%map accounts = accounts in
    List.map accounts ~f:(fun { pk; sk; balance; timing } ->
        let account_id = Account_id.create pk Token_id.default in
        let balance = Balance.of_mina_string_exn (Int.to_string balance) in
        (Some sk, account_with_timing account_id balance timing) )
end

(** Generate a ledger using the sample keypairs from [Mina_base] with the given
    balances.
*)
module Balances (Balances : Intf.Named_balances_intf) = struct
  open Intf.Private_accounts

  include Private_accounts (struct
    include Balances

    let accounts =
      let open Lazy.Let_syntax in
      let%map balances = Balances.balances
      and keypairs = Key_gen.Sample_keypairs.keypairs in
      List.mapi balances ~f:(fun i b ->
          { balance = b
          ; pk = fst keypairs.(i)
          ; sk = snd keypairs.(i)
          ; timing = Untimed
          } )
  end)
end

module Utils = struct
  (* Create a new [Ledger.Root.t] ledger with the components of a root
     backing_ledger for a genesis ledger. These components are the underlying
     root ledger and the stored mask on top of that root that is presented to
     users of the genesis root. The mask is passed in to this function so we can
     assert that there are no uncommitted changes to it, so we can simply
     checkpoint the root instead of performing a much slower account transfer. *)
  let create_root_from_backing_root genesis_mask root ~config ~depth () =
    assert (
      Ledger_hash.equal
        (Ledger.merkle_root genesis_mask)
        (Root_ledger.merkle_root root) ) ;
    assert (Root_ledger.depth root = depth) ;
    Root_ledger.create_checkpoint ~config root () |> Or_error.return

  let create_root_from_backing_root_with_directory genesis_mask root ~directory
      ~depth () =
    assert (
      Ledger_hash.equal
        (Ledger.merkle_root genesis_mask)
        (Root_ledger.merkle_root root) ) ;
    assert (Root_ledger.depth root = depth) ;
    Root_ledger.create_checkpoint_with_directory root ~directory_name:directory
    |> Or_error.return

  let keypair_of_account_record_exn (private_key, account) =
    let open Account in
    let sk_error_msg =
      "cannot access genesis ledger account private key "
      ^ "(HINT: did you forget to compile with `--profile=test`?)"
    in
    let pk_error_msg = "failed to decompress a genesis ledger public key" in
    let private_key = Option.value_exn private_key ~message:sk_error_msg in
    let public_key =
      Option.value_exn
        (Public_key.decompress account.public_key)
        ~message:pk_error_msg
    in
    { Keypair.public_key; private_key }

  let id_of_account_record (_private_key, account) = Account.identifier account

  let pk_of_account_record (_private_key, account) = Account.public_key account

  let find_account_record_exn ~f accounts =
    List.find_exn accounts ~f:(fun (_, account) -> f account)

  let find_new_account_record_exn_ accounts old_account_pks =
    find_account_record_exn accounts ~f:(fun new_account ->
        not
          (List.mem ~equal:Public_key.Compressed.equal old_account_pks
             (Account.public_key new_account) ) )

  let find_new_account_record_exn accounts old_account_pks =
    find_new_account_record_exn_ accounts
      (List.map ~f:Public_key.compress old_account_pks)
end

include Utils

module Make (Inputs : Intf.Ledger_input_intf) : Intf.S = struct
  include Inputs

  (* TODO: #1488 compute this at compile time instead of lazily *)
  (* The backing_ledger is either an `Ephemeral ledger, which exists purely as a
     mask in memory, with a null ledger root, or a `Root ledger, which is backed
     by a concrete [Root_ledger.t] on disk. We need a single mask to present to
     the users of the genesis ledger interface (through the [t] value below), so
     that is also saved here. *)
  let backing_ledger =
    let open Lazy.Let_syntax in
    let%map ledger, insert_accounts =
      match directory with
      | `Ephemeral ->
          lazy (`Ephemeral (Ledger.create_ephemeral ~depth ()), true)
      | `New backing_type ->
          lazy
            ( `Root
                (Root_ledger.create_temporary ~logger ~backing_type ~depth ())
            , true )
      | `Path (directory_name, backing_type) ->
          lazy
            ( `Root
                (Root_ledger.create ~logger
                   ~config:
                     (Root_ledger.Config.with_directory ~backing_type
                        ~directory_name )
                   ~depth () )
            , false )
    in
    let masked =
      match ledger with
      | `Ephemeral ledger ->
          ledger
      | `Root ledger ->
          Root_ledger.as_masked ledger
    in
    ( if insert_accounts then
      let addrs_and_accounts =
        let ledger_depth = Ledger.depth masked in
        Lazy.force accounts
        |> List.mapi ~f:(fun i (_, acct) ->
               (Ledger.Addr.of_int_exn ~ledger_depth i, acct) )
      in
      Ledger.set_batch_accounts masked addrs_and_accounts ) ;
    (ledger, masked)

  let t = Lazy.map ~f:snd backing_ledger

  let create_root ~config ~depth () =
    let backing_ledger, mask = Lazy.force backing_ledger in
    match backing_ledger with
    | `Ephemeral ledger ->
        let open Or_error.Let_syntax in
        let root = Root_ledger.create ~logger ~config ~depth () in
        (* We are transferring to an unmasked view of the root, so this is
           used solely for the transfer side effect *)
        let%map _dest =
          Ledger_transfer_mask.transfer_accounts ~src:ledger
            ~dest:(Root_ledger.as_unmasked root)
        in
        root
    | `Root ledger ->
        create_root_from_backing_root mask ledger ~config ~depth ()

  let create_root_with_directory ~directory ~depth () =
    let backing_ledger, mask = Lazy.force backing_ledger in
    match backing_ledger with
    | `Ephemeral ledger ->
        (* For ephemeral ledgers, create a fresh root with stable_db backing
           and transfer accounts to it *)
        let open Or_error.Let_syntax in
        let config =
          Root_ledger.Config.with_directory ~backing_type:Stable_db
            ~directory_name:directory
        in
        let root = Root_ledger.create ~logger ~config ~depth () in
        (* We are transferring to an unmasked view of the root, so this is
           used solely for the transfer side effect *)
        let%map _dest =
          Ledger_transfer_mask.transfer_accounts ~src:ledger
            ~dest:(Root_ledger.as_unmasked root)
        in
        root
    | `Root ledger ->
        create_root_from_backing_root_with_directory mask ledger ~directory
          ~depth ()

  include Utils

  let find_account_record_exn ~f =
    find_account_record_exn ~f (Lazy.force accounts)

  let find_new_account_record_exn_ old_account_pks =
    find_new_account_record_exn_ (Lazy.force accounts) old_account_pks

  let find_new_account_record_exn old_account_pks =
    find_new_account_record_exn (Lazy.force accounts) old_account_pks

  let largest_account_exn =
    let error_msg =
      "cannot calculate largest account in genesis ledger: "
      ^ "genesis ledger has no accounts"
    in
    Memo.unit (fun () ->
        List.max_elt (Lazy.force accounts) ~compare:(fun (_, a) (_, b) ->
            Balance.compare a.balance b.balance )
        |> Option.value_exn ?here:None ?error:None ~message:error_msg )

  let largest_account_id_exn =
    Memo.unit (fun () -> largest_account_exn () |> id_of_account_record)

  let largest_account_pk_exn =
    Memo.unit (fun () -> largest_account_exn () |> pk_of_account_record)

  let largest_account_keypair_exn =
    Memo.unit (fun () -> keypair_of_account_record_exn (largest_account_exn ()))
end

module Packed = struct
  type t = (module Intf.S)

  let t ((module L) : t) = L.t

  let create_root ((module L) : t) = L.create_root

  let create_root_with_directory ((module L) : t) = L.create_root_with_directory

  let depth ((module L) : t) = L.depth

  let accounts ((module L) : t) = L.accounts

  let find_account_record_exn ((module L) : t) = L.find_account_record_exn

  let find_new_account_record_exn_ ((module L) : t) =
    L.find_new_account_record_exn_

  let find_new_account_record_exn ((module L) : t) =
    L.find_new_account_record_exn

  let largest_account_exn ((module L) : t) = L.largest_account_exn ()

  let largest_account_id_exn ((module L) : t) = L.largest_account_id_exn ()

  let largest_account_pk_exn ((module L) : t) = L.largest_account_pk_exn ()

  let largest_account_keypair_exn ((module L) : t) =
    L.largest_account_keypair_exn ()
end

module Of_ledger (T : sig
  val backing_ledger : Root_ledger.t Lazy.t

  val depth : int
end) : Intf.S = struct
  include T

  let backing_ledger =
    Lazy.map
      ~f:(fun ledger -> (ledger, Root_ledger.as_masked ledger))
      backing_ledger

  let t = Lazy.map ~f:snd backing_ledger

  let accounts =
    Lazy.map t
      ~f:(Ledger.foldi ~init:[] ~f:(fun _loc accs acc -> (None, acc) :: accs))

  include Utils

  let create_root ~config ~depth () =
    let genesis_root, mask = Lazy.force backing_ledger in
    create_root_from_backing_root mask genesis_root ~config ~depth ()

  let create_root_with_directory ~directory ~depth () =
    let genesis_root, mask = Lazy.force backing_ledger in
    create_root_from_backing_root_with_directory mask genesis_root ~directory
      ~depth ()

  let find_account_record_exn ~f =
    find_account_record_exn ~f (Lazy.force accounts)

  let find_new_account_record_exn_ old_account_pks =
    find_new_account_record_exn_ (Lazy.force accounts) old_account_pks

  let find_new_account_record_exn old_account_pks =
    find_new_account_record_exn (Lazy.force accounts) old_account_pks

  let largest_account_exn =
    let error_msg =
      "cannot calculate largest account in genesis ledger: "
      ^ "genesis ledger has no accounts"
    in
    Memo.unit (fun () ->
        List.max_elt (Lazy.force accounts) ~compare:(fun (_, a) (_, b) ->
            Balance.compare a.Account.balance b.Account.balance )
        |> Option.value_exn ?here:None ?error:None ~message:error_msg )

  let largest_account_id_exn =
    Memo.unit (fun () -> largest_account_exn () |> id_of_account_record)

  let largest_account_pk_exn =
    Memo.unit (fun () -> largest_account_exn () |> pk_of_account_record)

  let largest_account_keypair_exn () =
    failwith "cannot access genesis ledger account private key"
end

let fetch_ledger, register_ledger =
  let ledgers = ref String.Map.empty in
  let register_ledger ((module Ledger : Intf.Named_accounts_intf) as l) =
    ledgers := Map.add_exn !ledgers ~key:Ledger.name ~data:l
  in
  let fetch_ledger name = Map.find !ledgers name in
  (fetch_ledger, register_ledger)

let fetch_ledger_exn name = Option.value_exn (fetch_ledger name)

module Register (Accounts : Intf.Named_accounts_intf) :
  Intf.Named_accounts_intf = struct
  let () = register_ledger (module Accounts)

  include Accounts
end

module Testnet_postake = Register (Balances (Testnet_postake_ledger))

module Testnet_postake_many_producers = Register (Balances (struct
  let name = "testnet_postake_many_producers"

  let balances =
    lazy
      (let high_balances = List.init 50 ~f:(Fn.const 5_000_000) in
       let low_balances = List.init 10 ~f:(Fn.const 1_000) in
       high_balances @ low_balances )
end))

module Test = Register (Balances (Test_ledger))

module Unit_test_ledger = Make (struct
  include Test

  let directory = `Ephemeral

  let depth =
    Genesis_constants.For_unit_tests.Constraint_constants.t.ledger_depth

  let logger = Logger.create ()
end)

let for_unit_tests : Packed.t = (module Unit_test_ledger)
